import numpy as np
import torch
import os, datetime

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def get_save_dir():
    base_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "deep-models", "strategy", "ppo-" + datetime.now().strftime("%Y%m%d%H%M%S"))
    os.makedirs(base_dir, exist_ok=True)
    return base_dir

class PPOMemory:
    def __init__(self, batch_size):
        self.states = []
        self.actions = []
        self.probs = []
        self.vals = []
        self.rewards = []
        self.dones = []
        self.batch_size = batch_size
    
    def store(self, state, action, prob, val, reward, done):
        self.states.append(state)
        self.actions.append(action)
        self.probs.append(float(prob))
        self.vals.append(float(val))
        self.rewards.append(reward)
        self.dones.append(done)
    
    def clear(self):
        self.states = []
        self.actions = []
        self.probs = []
        self.vals = []
        self.rewards = []
        self.dones = []
    
    def get_batches(self):
        n_states = len(self.states)
        batch_start = np.arange(0, n_states, self.batch_size)
        indices = np.arange(n_states, dtype=np.int64)
        np.random.shuffle(indices)
        batches = [indices[i:i+self.batch_size] for i in batch_start]
        
        return np.array(self.states), np.array(self.actions), np.array(self.probs), \
               np.array(self.vals), np.array(self.rewards), np.array(self.dones), batches
